from __future__ import annotations

from typing import Dict, Any
import json

from ..utils.ai_utils import AIClient
import logging


class ColumnTitleAgent:
    """Agent: Given final SQL, alias->expression map, brief column comments and the NL query,
    produce friendly Chinese titles for columns.

    The agent must return a JSON object mapping alias -> title.
    """

    def generate_titles(
        self,
        query: str,
        final_sql: str,
        alias_expr: Dict[str, str] | None,
        brief_comments: Dict[str, str] | None,
        col_names: list[str],
    ) -> Dict[str, str]:
        logger = logging.getLogger(__name__)
        alias_expr_json = json.dumps(alias_expr or {}, ensure_ascii=False)
        comments_json = json.dumps(brief_comments or {}, ensure_ascii=False)
        logger.info("ColumnTitleAgent brief_comments: %s", brief_comments)

        prompt = (
            "你是资深数据分析师，任务是为 SQL 查询结果的列生成用户可读的中文标题。\n"
            "规则：\n"
            "1) 聚合/计算列必须采用‘函数（字段备注）’格式：若该字段在提供的备注中存在，严格使用备注；若无备注，再用字段语义替代。例：SUM(amount)→求和（金额），COUNT(*)→计数（全部），MAX(discount)→最大（折扣）。\n"
            "2) 非聚合列直接用字段备注；无备注时根据字段名语义判断（amount/revenue→金额，price→价格，quantity/count→数量，discount→折扣，time/date→时间，name→名称 等）。\n"
            "3) 仅返回 JSON，键为列名（别名），值为中文标题，不要任何解释。若列名与表达式匹配不上，请根据表达式直接命名。\n"
            f"查询: {query}\n最终SQL: {final_sql}\n列名: {', '.join(col_names)}\n别名来源表达式(JSON): {alias_expr_json}\n相关字段备注(JSON): {comments_json}"
        )

        logger.info("ColumnTitleAgent request: cols=%s alias_expr=%s comments=%s sql=%s query=%s", col_names, alias_expr, brief_comments, final_sql, query)
        try:
            content = AIClient.chat(
                model=None,
                messages=[
                    {"role": "system", "content": "严格按用户要求，仅返回 JSON。"},
                    {"role": "user", "content": prompt},
                ],
                temperature=0.2,
                max_tokens=400,
            )
            logger.info("ColumnTitleAgent raw response: %s", content)
        except Exception:
            logger.exception("ColumnTitleAgent chat failed")
            return {}

        try:
            # 支持从代码块中提取
            import re  # noqa
            m = re.search(r"```[a-zA-Z]*\n([\s\S]*?)\n```", content)
            text = m.group(1).strip() if m else content.strip()
            obj = json.loads(text)
            if isinstance(obj, dict):
                logger.info("ColumnTitleAgent parsed mapping: %s", obj)
                return obj
            logger.info("ColumnTitleAgent parsed non-dict: %s", type(obj))
            return {}
        except Exception:
            logger.exception("ColumnTitleAgent parsing failed")
            return {}


